# -*- coding:utf-8 -*-
"""
使用wind提取股票的相关信息

从wind中提取的字段有：
'SEC_NAME', 股票简称
'PRENAME', 证券曾用名
'SEC_NAME1', 证券简称（支持历史）
'COMP_NAME', 公司中文名称
'CHAIRMAN', 法定代表人
'BOARDCHAIRMEN', 董事长
'CEO', 总经理
'holder_name',大股东名称
'holder_controller', 实际控制人名称
'holder_liqname', 流通股东名称

通过这些字段构建股票关系图谱
"""
from WindPy import *
import os
import pickle
from string import punctuation
import re

def removePunctuation(content):
    punc = punctuation + '.,;《》？！“”‘’@#￥%…&×（）—+【】{};；●，。·&～、|\:：'
    return re.sub('[{}]'.format(punc), '', content)

# 读取股票列表信息
stockList = []
with open('../data/stockList.csv', 'r', encoding='utf-8') as f:
    for each in f.readlines():
        stockList.append(each.strip())

# 读取windData的日期
date_wind = []
for year in range(2012, 2018):
    for date_str in ['0501', '0901', '1101']:
        date_wind.append(str(year)+date_str)

wind_data_path = '../data/windData'

# 将所有出现过的名词都放入词库中
financialDict = set()

for date in date_wind:
    # 创建属于该date的pkl文件夹
    pklPath = os.path.join('../data/pklFiles', date)
    os.makedirs(pklPath)
    windFile = open(os.path.join(wind_data_path, date+'.pkl'), 'rb')
    data = pickle.load(windFile)
    # 拿到详细的数据字段
    detailData = data.Data
    # 将所有字段构建关系映射
    # 1. 股票代码的不同前后缀变形
    stockCodeDiffVariation = {}
    for stock in stockList:
        stockCode = stock.strip()[:-3]
        financialDict.add(stockCode)
        # stockCodeDiffVariation[000001] = [xx,xxx]，将股票代码所有能够映射到的实体都保存下来
        stockCodeDiffVariation[stockCode] = []
        if stockCode[0] == '6':
            stockCodeDiffVariation[stockCode].append(stockCode+'SH')
            stockCodeDiffVariation[stockCode].append(stockCode+'Sh')
            stockCodeDiffVariation[stockCode].append(stockCode+'sH')
            stockCodeDiffVariation[stockCode].append(stockCode+'sh')
            stockCodeDiffVariation[stockCode].append('SH'+stockCode)
            stockCodeDiffVariation[stockCode].append('Sh'+stockCode)
            stockCodeDiffVariation[stockCode].append('sH'+stockCode)
            stockCodeDiffVariation[stockCode].append('sh'+stockCode)
        else:
            stockCodeDiffVariation[stockCode].append(stockCode+'SZ')
            stockCodeDiffVariation[stockCode].append(stockCode+'Sz')
            stockCodeDiffVariation[stockCode].append(stockCode+'sZ')
            stockCodeDiffVariation[stockCode].append(stockCode+'sz')
            stockCodeDiffVariation[stockCode].append('SZ'+stockCode)
            stockCodeDiffVariation[stockCode].append('Sz'+stockCode)
            stockCodeDiffVariation[stockCode].append('sZ'+stockCode)
            stockCodeDiffVariation[stockCode].append('sz'+stockCode)
        # stockCodeDiffVariation[xx] = 000001，所有的实体同样要映射到股票代码
        for obj in stockCodeDiffVariation[stockCode]:
            stockCodeDiffVariation[obj] = stockCode
            financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    stockCodeDiffVariationFile = open(os.path.join(pklPath, 'stockCodeDiffVariation.pkl'), 'wb')
    pickle.dump(stockCodeDiffVariation, stockCodeDiffVariationFile, -1)

    # 2.sec_name：股票简称
    secName = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        secName[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[0][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                element = removePunctuation(element)
                tempSet.add(element) # 原始信息
                # 万科(退市)
                if element[-4:] == '(退市)':
                    element = element[:-4]
                # 万科A
                if element[-1] == 'A':
                    element = element[:-1]
                # G万科
                if element[0] == 'G':
                    element = element[1:]
                # SST万科
                if element[:3] == 'SST':
                    element = element[3:]
                # ST万科
                if element[:2] == 'ST':
                    element = element[2:]
            for obj in tempSet:
                secName[stockCode].append(obj) # stockCode映射到实体
                secName[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    secNameFile = open(os.path.join(pklPath, 'secName.pkl'), 'wb')
    pickle.dump(secName, secNameFile, -1)

    # 3.PRENAME：证券曾用名
    preName = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        preName[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[1][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                element = removePunctuation(element)
                tempSet.add(element) # 原始信息
                # 万科(退市)
                if element[-4:] == '(退市)':
                    element = element[:-4]
                # 万科A
                if element[-1] == 'A':
                    element = element[:-1]
                # G万科
                if element[0] == 'G':
                    element = element[1:]
                # SST万科
                if element[:3] == 'SST':
                    element = element[3:]
                # ST万科
                if element[:2] == 'ST':
                    element = element[2:]
            for obj in tempSet:
                preName[stockCode].append(obj) # stockCode映射到实体
                preName[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    preNameFile = open(os.path.join(pklPath, 'preName.pkl'), 'wb')
    pickle.dump(preName, preNameFile, -1)

    # 4.sec_name1：证券简称（支持历史）
    secName1 = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        secName1[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[2][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                element = removePunctuation(element)
                tempSet.add(element) # 原始信息
                # 万科(退市)
                if element[-4:] == '(退市)':
                    element = element[:-4]
                # 万科A
                if element[-1] == 'A':
                    element = element[:-1]
                # G万科
                if element[0] == 'G':
                    element = element[1:]
                # SST万科
                if element[:3] == 'SST':
                    element = element[3:]
                # ST万科
                if element[:2] == 'ST':
                    element = element[2:]
            for obj in tempSet:
                secName1[stockCode].append(obj) # stockCode映射到实体
                secName1[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    secName1File = open(os.path.join(pklPath, 'secName1.pkl'), 'wb')
    pickle.dump(secName1, secName1File, -1)

    # 5. COMP_NAME: 公司中文名称
    compName = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        compName[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[3][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                element = removePunctuation(element)
                tempSet.add(element) # 原始信息
                # 万科(退市)
                if element[-4:] == '(退市)':
                    element = element[:-4]
                # 万科A
                if element[-1] == 'A':
                    element = element[:-1]
                # G万科
                if element[0] == 'G':
                    element = element[1:]
                # SST万科
                if element[:3] == 'SST':
                    element = element[3:]
                # ST万科
                if element[:2] == 'ST':
                    element = element[2:]
            for obj in tempSet:
                compName[stockCode].append(obj) # stockCode映射到实体
                compName[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    compNameFile = open(os.path.join(pklPath, 'compName.pkl'), 'wb')
    pickle.dump(compName, compNameFile, -1)

    # 6. CHAIRMAN: 法定代表人
    chairman = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        chairman[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[4][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                for e in element.split(','):
                    e = removePunctuation(e)
                    tempSet.add(e)
            for obj in tempSet:
                chairman[stockCode].append(obj) # stockCode映射到实体
                chairman[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    chairmanFile = open(os.path.join(pklPath, 'chairman.pkl'), 'wb')
    pickle.dump(chairman, chairmanFile, -1)

    # 7. BOARDCHAIRMEN: 董事长
    boardchairman = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        boardchairman[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[5][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                for e in element.split(','):
                    e = removePunctuation(e)
                    tempSet.add(e)
            for obj in tempSet:
                boardchairman[stockCode].append(obj) # stockCode映射到实体
                boardchairman[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    boardchairmanFile = open(os.path.join(pklPath, 'boardchairman.pkl'), 'wb')
    pickle.dump(boardchairman, boardchairmanFile, -1)

    # 8. CEO: 总经理
    ceo = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        ceo[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[6][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                for e in element.split(','):
                    e = removePunctuation(e)
                    tempSet.add(e)
            for obj in tempSet:
                ceo[stockCode].append(obj) # stockCode映射到实体
                ceo[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    ceoFile = open(os.path.join(pklPath, 'ceo.pkl'), 'wb')
    pickle.dump(ceo, ceoFile, -1)

    # 9. holder_name: 大股东名称
    holderName = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        holderName[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[7][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                for e in element.split(','):
                    e = removePunctuation(e)
                    tempSet.add(e)
                    # 对于其中的一些产品名字，将“公司”后面进行截断操作
                    index = e.rfind('公司')
                    if index != -1:
                        tempSet.add(e[:index+2])
            for obj in tempSet:
                holderName[stockCode].append(obj) # stockCode映射到实体
                holderName[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    holderNameFile = open(os.path.join(pklPath, 'holderName.pkl'), 'wb')
    pickle.dump(holderName, holderNameFile, -1)

    # 10. holder_controller: 实际控制人名称
    holderController = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        holderController[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[8][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                for e in element.split(','):
                    e = removePunctuation(e)
                    tempSet.add(e)
                    # 对于其中的一些产品名字，将“公司”后面进行截断操作
                    index = e.rfind('公司')
                    if index != -1:
                        tempSet.add(e[:index + 2])
            for obj in tempSet:
                holderController[stockCode].append(obj) # stockCode映射到实体
                holderController[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    holderControllerFile = open(os.path.join(pklPath, 'holderController.pkl'), 'wb')
    pickle.dump(holderController, holderControllerFile, -1)

    # 11. holder_liqname: 流通股东名称
    holderLiqname = {}
    for stockIndex in range(len(stockList)):
        stock = stockList[stockIndex]
        stockCode = stock[:-3]
        holderLiqname[stockCode] = [] # 保存所有stockCode能够映射到的实体
        tempSet = set() # 临时用来过滤重复的实体
        content = detailData[9][stockIndex]
        if content != None:
            content = content.strip().split(';')
            for element in content:
                element = element.strip()
                for e in element.split(','):
                    e = removePunctuation(e)
                    tempSet.add(e)
                    # 对于其中的一些产品名字，将“公司”后面进行截断操作
                    index = e.rfind('公司')
                    if index != -1:
                        tempSet.add(e[:index + 2])
            for obj in tempSet:
                holderLiqname[stockCode].append(obj) # stockCode映射到实体
                holderLiqname[obj] = stockCode # 实体映射到stockCode
                financialDict.add(obj) # 将所有的实体放入字典
    # 将数据进行序列化
    holderLiqnameFile = open(os.path.join(pklPath, 'holderLiqname.pkl'), 'wb')
    pickle.dump(holderLiqname, holderLiqnameFile, -1)

# 融合financialDict.csv
with open('../dictionary/financialDict.csv', 'r', encoding='utf-8') as f,\
    open('../dictionary/biggerFinancialDict.csv', 'w', encoding='utf-8')as o:
    for each in f.readlines():
        financialDict.add(each.strip())
    output = sorted(list(financialDict))
    for each in output:
        o.write("%s\n" % each)
